# src/cli.py

import os
import sys
import logging
import argparse
import traceback
import concurrent.futures

from src.ctrl import (
    request_stop as ctrl_request_stop,
    reset_stop as ctrl_reset_stop,
    is_stopped as ctrl_is_stopped,
)

from src.tcrun import TcRunner
from src.exphd import ExpHandler


logging.basicConfig(level=logging.INFO, format="%(asctime)s[%(levelname)s]%(message)s")
logger = logging.getLogger(__name__)


def _process_dir_worker(dir: str):
    """
    Module-level top function for subprocess execution of work dir processing (can be pickled).

    Args:
        dir: Work dir to pro
    Returns:
        Tuple: (dir, err, trace)

    """
    # Resources needs to be re-instantiated in subprocess
    try:
        exphdlr = ExpHandler()
        exphdlr.exps2xlsx(dir)
        return (dir, None, None)
    except Exception as e:
        return (dir, str(e), traceback.format_exc())


def process_dirs(dirs: list[str], parallel: int = 2):
    """
    Process dirs in a process pool, with each dir handled by a separate process.

    Args:
        dirs: List of dir to process
        parallel: Number of processes to use for parallel processing
    """
    if len(dirs) <= 0:
        logger.info("No dirs found to process")
        return
    logger.info(f"{parallel} processes to process {len(dirs)} dirs")

    with concurrent.futures.ProcessPoolExecutor(max_workers=parallel) as exe:
        futures = {exe.submit(_process_dir_worker, d): d for d in dirs}

        try:
            for fut in concurrent.futures.as_completed(futures):
                if ctrl_is_stopped():
                    logger.info("Interrupt signal received, stopping tasks")
                    break
                dir = futures.get(fut)
                dir, err, trace = fut.result()
                if err:
                    logger.error(f"Error processing {dir}: {err}")
                    logger.debug(trace)
                else:
                    logger.info(f"Subprocess processed {dir} successfully")
        except KeyboardInterrupt:
            logger.info("Keyboard interrupt, stopping process pool")
        except Exception as e:
            logger.error(f"Task exception: {e}")
        finally:
            logger.info("Processing completed")


def request_stop():
    """Request to stop processing (wrapper)."""
    ctrl_request_stop()
    logger.debug("Interrupt signal sent")


def reset_stop():
    """Reset interrupt flag (wrapper)."""
    ctrl_reset_stop()


def run(config_path: str, parallel: int, only_exps: str):
    """
    Program entry.

    Args:
        config_path: Configuration file path
        async_mode: Whether to use asynchronous mode
        only_exps: Process only the specified dir's EXP files
    """
    # Reset interrupt flag
    reset_stop()

    if only_exps:
        logger.info(f"Process only EXP files (parallel, {parallel} processes)...")
        # Collect all dirs to be processed
        work_dirs = []
        for root, dirs, files in os.walk(only_exps):
            if ctrl_is_stopped():
                logger.info("Interrupt signal received, stopping collecting dirs")
                break
            work_dirs.append(root)

    else:
        logger.info("Calculation and processing mode running...")
        runner = TcRunner()
        work_dirs = [d for d in runner.run(config_path, parallel=parallel)]

    process_dirs(work_dirs, parallel=parallel)


def cli():
    parser = argparse.ArgumentParser(
        description="TcExpProc - Thermodynamic Calculation Experiment Processor",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run with configuration file
  %(prog)s --config tcexp_cfg.toml --parallel 1

  # Process only EXP files in specified dir
  %(prog)s --exps-dir path/to/exp/files --parallel 1

  # Run in debug mode
  %(prog)s --config tcexp_cfg.toml --mode sync --log-level debug
""",
    )

    # Processing mode option group (config file/EXP files)
    source_group = parser.add_mutually_exclusive_group(required=True)
    source_group.add_argument(
        "--config",
        type=str,
        metavar="FILE",
        help="Run with configuration file, specify TOML configuration file path",
    )
    source_group.add_argument(
        "--exps-dir",
        type=str,
        dest="only_exps",
        metavar="DIR",
        help="Process only EXP files in the specified dir, without calculation",
    )

    # Parallel count (for parallel processing EXP or tasks), default 1 means no parallel
    parser.add_argument(
        "--parallel",
        type=int,
        default=1,
        help="Parallel count: >1 means parallel processing (default: 1, serial)",
    )

    # Log level option
    parser.add_argument(
        "--log-level",
        choices=["debug", "info"],
        default="info",
        help="Log output level: debug=debug information, info=general information (default: info)",
    )

    args = parser.parse_args()

    # Set log level
    logger.setLevel(logging.DEBUG if args.log_level == "debug" else logging.INFO)

    # Run program
    run(
        config_path=args.config
        or os.path.join(os.path.dirname(sys.argv[0]), "tcexp_cfg.toml"),
        parallel=args.parallel,
        only_exps=args.only_exps,
    )


if __name__ == "__main__":
    cli()
